
"""
子任务分解模块
"""
from core.data_analyze.state import AnalyzeState
from core.data_analyze.prompts import data_analyze_prompts
from langchain_core.messages import AIMessage
from langgraph.config import get_stream_writer
from langgraph.types import Command
from json_repair import repair_json
from typing import Any
from core.llm import llm
import pandas as pd
from loguru import logger
import json


def task_decompose_node(state: AnalyzeState):

    writer = get_stream_writer()
    writer("## **数据分析任务分解**")
    logger.info("## **数据分析任务分解**")

    datasets = state.get('datasets')
    table_preview = []
    for dataset in datasets:
        table_preview.append(f"** {dataset.get('description')} **\n")
        df = pd.DataFrame(data=dataset.get('rows'))
        sample_count = min(20, df.shape[0])
        table_preview.append(df.sample(sample_count).to_markdown())
    table_preview = '\n'.join(table_preview)
    writer("**数据预览：**")
    logger.info("**数据预览：**")
    # writer(table_preview)
    # logger.info(table_preview)

    template_prompt = data_analyze_prompts.get('task_decompose')
    prompt = template_prompt.format(
        inputs={
            "task_description": state.get('revised_task'),
            "table_preview": table_preview
        },
        remove_template_variables=True
    )

    max_retry_times = state.get('max_retry_times')
    cur_retry_times = state.get('cur_retry_times')
    while cur_retry_times != max_retry_times:
        content = llm.invoke(prompt).content
        decomposed_tasks = repair_json(content, return_objects=True)
        writer("**分解的pandas任务**")
        logger.info("**分解的pandas任务**")
        writer(decomposed_tasks)
        logger.info(json.dumps(decomposed_tasks, ensure_ascii=False, indent=4))
        if isinstance(decomposed_tasks, list):
            decomposed_tasks = list(filter(validate_task, decomposed_tasks))
            if len(decomposed_tasks) != 0:
                return Command(
                    update={
                        "messages": [AIMessage(content=decomposed_tasks, name="task_decompose")],
                        "decomposed_tasks": decomposed_tasks
                    },
                    goto="data_analyze",
                )
        cur_retry_times += 1

    return Command(goto="__end__")


def validate_task(task: Any) -> bool:
    """验证分解的任务是否有效"""
    if not isinstance(task, dict):
        return False
    task_name = task.get('task_name', '')
    if task_name is None or task_name == '':
        return False
    if task.get('out_type', 'string') not in ['string', 'dataframe', 'json', 'plot']:
        return False
    return True